import os
import requests
import gradio as gr
from datetime import datetime
from PIL import Image
from io import BytesIO
import time

# 定义尺寸映射字典
size_mapping = {
    "9:16（竖屏）": "1536x2730",
    "16:9（横屏）": "2730x1536",
    "1:1（正方形）": "2048x2048",
    "4:3（传统显示器）": "2300x1800",
    "3:4（竖向4:3）": "1800x2300",
    "2:3（宽屏）": "1600x2400",
    "3:2（35mm胶片）": "2400x1600",
    "21:9（超宽屏）": "2520x1080",
}

# 1536 定义尺寸映射字典
# size_mapping = {
#     "9:16（竖屏）": "1152x2048",
#     "16:9（横屏）": "2048x1152",
#     "1:1（正方形）": "1536x1536",
#     "4:3（传统显示器）": "1772x1329",
#     "3:4（竖向4:3）": "1329x1772",
#     "2:3（宽屏）": "2352x1008",
#     "3:2（35mm胶片）": "1881x1254",
#     "21:9（超宽屏）": "2352x1008",
# }

# 获取保存路径和API Key（从环境变量）
save_dir = os.environ.get("BOZO") + "/Downloads/"
api_key = os.environ.get("ARK_API_KEY")

# 确保保存目录存在
os.makedirs(save_dir, exist_ok=True)

def generate_image(prompt, size_option):
    if not prompt:
        return None, "请输入提示词", "请输入提示词"
    
    if not api_key:
        return None, "API密钥未配置", "API密钥未配置"
    
    try:
        # 准备API请求
        url = "https://ark.cn-beijing.volces.com/api/v3/images/generations"
        selected_size = size_mapping[size_option]
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        
        data = {
            "model": "doubao-seedream-3-0-t2i-250415",
            "prompt": prompt,
            "response_format": "url",
            "size": selected_size,
            "seed": -1,
            "guidance_scale": 2.5,
            "watermark": False
        }
        
        # 发送请求
        start_time = time.time()
        response = requests.post(url, headers=headers, json=data)
        end_time = time.time()
        
        if response.status_code == 200:
            response_data = response.json()
            image_url = response_data["data"][0]["url"]
            
            # 下载并保存图片
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"豆包_{timestamp}.jpeg"
            save_path = os.path.join(save_dir, filename)
            
            image = Image.open(BytesIO(requests.get(image_url).content))
            image.save(save_path)
            
            return image, image_url, f"图片生成成功！耗时: {end_time - start_time:.2f} 秒"
        else:
            return None, f"请求失败", f"状态码：{response.status_code}，错误信息：{response.text}"
    
    except Exception as e:
        return None, "生成失败", f"错误信息：{str(e)}"

# 创建Gradio界面
with gr.Blocks(title="豆包图片生成器") as demo:
    gr.Markdown("# 豆包图片生成器")
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="提示词", placeholder="请输入中文提示词...", lines=30)
            size = gr.Dropdown(
                label="图片尺寸",
                choices=list(size_mapping.keys()),
                value="9:16（竖屏）"
            )
            generate_btn = gr.Button("生成图片", variant="primary")
            url_output = gr.Textbox(label="图片URL地址")
            status_output = gr.Textbox(label="状态信息")
        
        with gr.Column(scale=2):
            image_output = gr.Image(label="生成的图片")
    
    # 设置事件
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, size],
        outputs=[image_output, url_output, status_output]
    )

# 启动服务，监听8001端口，允许外部访问
if __name__ == "__main__":
    demo.launch(server_port=8001, server_name="0.0.0.0")
